#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Dec 30 2021
Comparing Hg and O3 dry deposition velocities for Amazon when adjusting resistances 
@author: arifeinberg
"""

#%%
#import os
#os.chdir('/Users/arifeinberg/target2/fs03/d0/arifein/python/')

import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from scipy import stats
from functions_Hg0 import get_mod_data, ds_sel_yr, annual_avg
import cartopy.crs as ccrs
import cartopy.feature as cf
from matplotlib import colors
import scipy.io as sio
import sys
#%% DV of observations
obs_DV_O3 = [0.88, 0.66, 0.85]
obs_DV_Hg = [0.145027882, 0.327105329, 0.313361408, 0.225400311, 0.164927057, 0.305115055, 0.109951371, 0.151183135, 0.167675841, 0.156680704]
#%% DV of model
data_mat = sio.loadmat('misc_Data/LT_AMZN_O3_SO2_orig_v2.mat')
O3_model_orig = data_mat['DV_lt'].squeeze()[1:]
data_mat = sio.loadmat('misc_Data/LT_AMZN_O3_SO2_IRI_65_v2.mat')
O3_model_IRI_65 = data_mat['DV_lt'].squeeze()[1:]
data_mat = sio.loadmat('misc_Data/LT_AMZN_O3_SO2_IRLU_500_IRI_100_v2.mat')
O3_model_IRI_100_IRLU_500 = data_mat['DV_lt'].squeeze()[1:]
data_mat = sio.loadmat('misc_Data/LT_AMZN_O3_SO2_IRI_1_v2.mat')
O3_model_IRI_1 = data_mat['DV_lt'].squeeze()[1:]

all_O3 = [O3_model_orig, O3_model_IRI_65, O3_model_IRI_100_IRLU_500, O3_model_IRI_1, obs_DV_O3] 

data_mat = sio.loadmat('misc_Data/SA_dep_F0_3e-05_v4_orig.mat')
Hg_model_orig = np.mean(data_mat['DV_lt'].squeeze(),1)
data_mat = sio.loadmat('misc_Data/SA_dep_F0_3e-05_v4_IRI_65.mat')
Hg_model_IRI_65 = np.mean(data_mat['DV_lt'].squeeze(),1)
data_mat = sio.loadmat('misc_Data/SA_dep_F0_3e-05_v4_IRLU_500_IRI_100.mat')
Hg_model_IRI_100_IRLU_500 = np.mean(data_mat['DV_lt'].squeeze(),1)
data_mat = sio.loadmat('misc_Data/SA_dep_F0_3e-05_v4_IRI_1.mat')
Hg_model_IRI_1 = np.mean(data_mat['DV_lt'].squeeze(),1)

all_Hg = [Hg_model_orig, Hg_model_IRI_65, Hg_model_IRI_100_IRLU_500, Hg_model_IRI_1, obs_DV_Hg]

#%% Create plot comparing model and measurement
f1,  axes1 = plt.subplots(1, 2, figsize=[13,5],
                          gridspec_kw=dict(hspace=0.3, wspace=0.4))
plt.subplots_adjust(right=0.65)

axes1 = axes1.flatten()
tests = ["Original (RI=200, RLU=1000)","RI=65","RI=100, RLU=500","RI=1", "Observed"]
list_colours = ['#1b9e77','#d95f02','#7570b3','#e7298a', 'k']
n_sim = len(all_O3) - 1

x_values = np.linspace(1.6/n_sim,  1 - 2/n_sim, n_sim) # for spacing plot
diff_x = (x_values[1] - x_values[0])*3. # for spacing plot

for i in range(n_sim):
    # calculate mean, and error bars
    mean_val = np.mean(all_O3[i])
    min_val = min(all_O3[i])
    max_val = max(all_O3[i])
    
    yerr_low = mean_val - min_val
    yerr_high = max_val - mean_val
    yerr_com = [[yerr_low],[yerr_high]]

    axes1[0].errorbar(x_values[i],mean_val, yerr=yerr_com,
                     fmt='d',color='k',markerfacecolor=list_colours[i],
                     capsize=4, lw=2, ms=10, label=tests[i] )
    print(mean_val)
mean_val = np.mean(all_O3[-1])
min_val = min(all_O3[-1])
max_val = max(all_O3[-1])
yerr_low = mean_val - min_val
yerr_high = max_val - mean_val
yerr_com = [[yerr_low],[yerr_high]]
print(mean_val)
axes1[0].errorbar(x_values[-1] + diff_x,mean_val, yerr=yerr_com,
                     fmt='o',color='k',markerfacecolor=list_colours[-1],
                     capsize=4, lw=2, ms=9, label=tests[-1])
axes1[0].set_xticks([])
axes1[0].set_ylabel('Dry deposition velocity (cm s$^{-1}$)', fontsize=18)
axes1[0].set_title('Ozone', fontsize=20)
axes1[0].set_xlim(0.35,0.65)
axes1[0].tick_params(axis='y', which='major', labelsize=15)

for i in range(n_sim):
    # calculate mean, and error bars
    mean_val = np.mean(all_Hg[i])
    min_val = min(all_Hg[i])
    max_val = max(all_Hg[i])
    
    yerr_low = mean_val - min_val
    yerr_high = max_val - mean_val
    yerr_com = [[yerr_low],[yerr_high]]

    axes1[1].errorbar(x_values[i],mean_val, yerr=yerr_com,
                     fmt='d',color='k',markerfacecolor=list_colours[i],
                     capsize=4, lw=2, ms=10)
    print(mean_val)
mean_val = np.mean(all_Hg[-1])
min_val = min(all_Hg[-1])
max_val = max(all_Hg[-1])
yerr_low = mean_val - min_val
yerr_high = max_val - mean_val
yerr_com = [[yerr_low],[yerr_high]]
print(mean_val)
axes1[1].errorbar(x_values[-1] + diff_x,mean_val, yerr=yerr_com,
                     fmt='o',color='k',markerfacecolor=list_colours[-1],
                     capsize=4, lw=2, ms=9)
axes1[1].set_xticks([])
axes1[1].set_ylabel('Dry deposition velocity (cm s$^{-1}$)', fontsize=18)
axes1[1].set_title('Gaseous Elemental Mercury', fontsize=20)
axes1[1].set_xlim(0.35,0.65)
axes1[1].tick_params(axis='y', which='major', labelsize=15)

f1.legend(loc='center right',  fontsize=16)
f1.savefig('Figures/SI_Hg0_O3_resist.pdf',bbox_inches = 'tight')

